import torch
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from model import ImageClassifier
from utils import load_model
import random

# CIFAR10类别名称
classes = ('plane', 'car', 'bird', 'cat', 'deer', 
           'dog', 'frog', 'horse', 'ship', 'truck')

def visualize_predictions(model_path='improved_cifar10_classifier.pth', num_images=6, seed=42):
    """
    可视化模型预测结果
    
    参数:
        model_path: 模型保存路径
        num_images: 要可视化的图像数量(必须是平方数，如1,4,9,16等)
        seed: 随机种子，确保可重复性
    """
    # 设置随机种子
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 加载模型
    model = ImageClassifier().to(device)
    load_model(model, model_path)
    model.eval()
    print("Model loaded successfully")

    # 加载测试数据(不进行数据增强)
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
    ])
    test_dataset = CIFAR10(root='./data', train=False, download=True, transform=transform)
    
    # 随机选择一些图片
    indices = random.sample(range(len(test_dataset)), num_images)
    images = []
    labels = []
    for idx in indices:
        img, label = test_dataset[idx]
        images.append(img)
        labels.append(label)
    
    # 转换为张量
    images = torch.stack(images).to(device)
    labels = torch.tensor(labels).to(device)

    # 获取预测结果
    with torch.no_grad():
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        probabilities = torch.nn.functional.softmax(outputs, dim=1)
        top_probs, top_classes = torch.topk(probabilities, 3, dim=1)

    # 创建可视化
    plt.figure(figsize=(12, 12))
    rows = int(np.sqrt(num_images))
    cols = int(np.ceil(num_images / rows))
    
    for i in range(num_images):
        plt.subplot(rows, cols, i + 1)
        
        # 反归一化并转换为numpy数组
        img = images[i].cpu().numpy().transpose((1, 2, 0))
        img = img * np.array([0.2470, 0.2435, 0.2616]) + np.array([0.4914, 0.4822, 0.4465])
        img = np.clip(img, 0, 1)
        
        plt.imshow(img)
        
        # 构建标题文本
        title = f"True: {classes[labels[i]]}\n"
        title += f"Pred: {classes[predicted[i]]} ({top_probs[i][0]:.2f})\n"
        title += "Top 3:\n"
        for j in range(3):
            title += f"{classes[top_classes[i][j]]}: {top_probs[i][j]:.2f}\n"
        
        # 根据预测是否正确设置颜色
        color = 'green' if predicted[i] == labels[i] else 'red'
        plt.title(title, color=color)
        plt.axis('off')
    
    plt.tight_layout()
    plt.savefig('model_predictions.png', bbox_inches='tight', dpi=300)
    plt.show()
    print("Predictions visualization saved to 'model_predictions.png'")

if __name__ == "__main__":
    visualize_predictions()